from ultralytics import YOLO
import cv2
import numpy as np

# 1. 加载模型
model_path = 'data/weights/best.pt'  # 替换为你的模型文件路径
model = YOLO(model_path)

# 2. 加载视频
video_path = "C:/Users/永生理想/Desktop/4.mp4"  # 替换为你的视频文件路径
cap = cv2.VideoCapture(video_path)

# 获取视频的原始属性
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)

# 3. 设置视频写入器
output_path = "C:/Users/永生理想/Desktop/output_video.mp4"  # 输出视频路径
fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 使用 mp4v 编码器
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

# 设置显示窗口为可调整大小
cv2.namedWindow('YOLOv8 Detection', cv2.WINDOW_NORMAL)

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # 4. 模型预测
    results = model(frame)

    # 5. 处理和可视化检测结果
    for result in results:
        boxes = result.boxes  # 检测框
        for box in boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])  # 坐标
            conf = box.conf[0]  # 置信度
            cls = int(box.cls[0])  # 类别
            label = f'{model.names[cls]} {conf:.2f}'

            # 绘制检测框和标签
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    # 6. 显示处理后的帧
    cv2.imshow('YOLOv8 Detection', frame)
    
    # 7. 写入帧到输出视频
    out.write(frame)

    # 按'q'键退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放资源
cap.release()
out.release()  # 释放视频写入器
cv2.destroyAllWindows()